{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Copy Task Using NTM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will see, how to perform copy tasks using NTM. The goal of the copy task is to see how NTM stores and recall the sequence of arbitrary length. We will feed the network a random sequence along with a marker indicating the end of a sequence. It has to learn to output the given input sequence. So, the network will store the input sequence to the memory and then it will read back from the memory. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we will import all the necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "import argparse\n",
    "from PIL import Image\n",
    "from PIL import ImageOps\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will define a class NTMCell where we will implement our neural turing machine. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NTMCell():\n",
    "    def __init__(self, rnn_size, memory_size, memory_vector_dim, read_head_num, write_head_num,\n",
    "                 addressing_mode='content_and_location', shift_range=1, reuse=False, output_dim=None):\n",
    "        \n",
    "        #initialize all the variables\n",
    "        self.rnn_size = rnn_size\n",
    "        self.memory_size = memory_size\n",
    "        self.memory_vector_dim = memory_vector_dim\n",
    "        self.read_head_num = read_head_num\n",
    "        self.write_head_num = write_head_num\n",
    "        self.addressing_mode = addressing_mode\n",
    "        self.reuse = reuse\n",
    "        self.step = 0\n",
    "        self.output_dim = output_dim\n",
    "        self.shift_range = shift_range\n",
    "        \n",
    "        #initialize controller as the basic rnn cell\n",
    "        self.controller = tf.nn.rnn_cell.BasicRNNCell(self.rnn_size)\n",
    "\n",
    "        \n",
    "        \n",
    "    def __call__(self, x, prev_state):\n",
    "        \n",
    "        \n",
    "        prev_read_vector_list = prev_state['read_vector_list']  \n",
    "        prev_controller_state = prev_state['controller_state']     \n",
    "    \n",
    "        controller_input = tf.concat([x] + prev_read_vector_list, axis=1)\n",
    "        \n",
    "        #next we pass the controller which is the RNN cell, the controller_input and prev_controller_state\n",
    "        with tf.variable_scope('controller', reuse=self.reuse):\n",
    "            controller_output, controller_state = self.controller(controller_input, prev_controller_state)\n",
    "\n",
    "        #initialize read and write heads    \n",
    "        num_parameters_per_head = self.memory_vector_dim + 1 + 1 + (self.shift_range * 2 + 1) + 1\n",
    "        num_heads = self.read_head_num + self.write_head_num\n",
    "        \n",
    "        total_parameter_num = num_parameters_per_head * num_heads + self.memory_vector_dim * 2 * self.write_head_num\n",
    "        \n",
    "        #initialize weight matrix and bias and compute the parameters\n",
    "        with tf.variable_scope(\"o2p\", reuse=(self.step > 0) or self.reuse):\n",
    "            o2p_w = tf.get_variable('o2p_w', [controller_output.get_shape()[1], total_parameter_num],\n",
    "                                    initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))\n",
    "            o2p_b = tf.get_variable('o2p_b', [total_parameter_num],\n",
    "                                    initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))\n",
    "            parameters = tf.nn.xw_plus_b(controller_output, o2p_w, o2p_b)\n",
    "            \n",
    "            \n",
    "        head_parameter_list = tf.split(parameters[:, :num_parameters_per_head * num_heads], num_heads, axis=1)\n",
    "        \n",
    "        erase_add_list = tf.split(parameters[:, num_parameters_per_head * num_heads:], 2 * self.write_head_num, axis=1)\n",
    "\n",
    "        \n",
    "        #previous write weight vector\n",
    "        prev_w_list = prev_state['w_list'] \n",
    "        \n",
    "        #previous Memory\n",
    "        prev_M = prev_state['M']\n",
    "        \n",
    "        w_list = []\n",
    "        p_list = []\n",
    "        \n",
    "        \n",
    "        #now, we will initialize some of the important parameters that we use for addressing. \n",
    "        for i, head_parameter in enumerate(head_parameter_list):\n",
    "\n",
    "            #key vector\n",
    "            k = tf.tanh(head_parameter[:, 0:self.memory_vector_dim])\n",
    "            \n",
    "            #key strength(beta)\n",
    "            beta = tf.sigmoid(head_parameter[:, self.memory_vector_dim]) * 10  \n",
    "            \n",
    "            #interpolation gate\n",
    "            g = tf.sigmoid(head_parameter[:, self.memory_vector_dim + 1])\n",
    "            \n",
    "            #shift matrix  \n",
    "            s = tf.nn.softmax(\n",
    "                head_parameter[:, self.memory_vector_dim + 2:self.memory_vector_dim + 2 + (self.shift_range * 2 + 1)]\n",
    "            )\n",
    "            \n",
    "            #sharpening factor\n",
    "            gamma = tf.log(tf.exp(head_parameter[:, -1]) + 1) + 1\n",
    "            \n",
    "            with tf.variable_scope('addressing_head_%d' % i):\n",
    "                w = self.addressing(k, beta, g, s, gamma, prev_M, prev_w_list[i]) \n",
    "\n",
    "            w_list.append(w)\n",
    "            p_list.append({'k': k, 'beta': beta, 'g': g, 's': s, 'gamma': gamma})\n",
    "\n",
    "        \n",
    "        \n",
    "        #We basically perform two important operations in NTM one is read and other is write operation.\n",
    "    \n",
    "        \n",
    "        #A. Read Operation:\n",
    "            # read operation is the linear combination  of weights and memory. \n",
    "        \n",
    "        #select the head to read from\n",
    "        read_w_list = w_list[:self.read_head_num]\n",
    "        \n",
    "        \n",
    "        read_vector_list = []\n",
    "        for i in range(self.read_head_num):  \n",
    "            #linear combination  of the weights and memory\n",
    "            read_vector = tf.reduce_sum(tf.expand_dims(read_w_list[i], dim=2) * prev_M, axis=1)\n",
    "            read_vector_list.append(read_vector)\n",
    "            \n",
    "\n",
    "        #B. Write Operation\n",
    "        \n",
    "            #Unlike read operation, write operation consists of two steps - erase and add. \n",
    "        \n",
    "        #select the head to write\n",
    "        write_w_list = w_list[self.read_head_num:]\n",
    "        M = prev_M\n",
    "        for i in range(self.write_head_num):\n",
    "            \n",
    "            #the erase vector will be multipled with weight vector to denote which location to erase \n",
    "            w = tf.expand_dims(write_w_list[i], axis=2)\n",
    "            erase_vector = tf.expand_dims(tf.sigmoid(erase_add_list[i * 2]), axis=1)\n",
    "            \n",
    "            #next we perform the add operation\n",
    "            add_vector = tf.expand_dims(tf.tanh(erase_add_list[i * 2 + 1]), axis=1)\n",
    "            M = M * (tf.ones(M.get_shape()) - tf.matmul(w, erase_vector)) + tf.matmul(w, add_vector)\n",
    "            \n",
    "    \n",
    "        \n",
    "        #controller output\n",
    "        if not self.output_dim:\n",
    "            output_dim = x.get_shape()[1]\n",
    "        else:\n",
    "            output_dim = self.output_dim\n",
    "            \n",
    "        with tf.variable_scope(\"o2o\", reuse=(self.step > 0) or self.reuse):\n",
    "            o2o_w = tf.get_variable('o2o_w', [controller_output.get_shape()[1], output_dim],\n",
    "                                    initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))\n",
    "            o2o_b = tf.get_variable('o2o_b', [output_dim],\n",
    "                                    initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))\n",
    "            NTM_output = tf.nn.xw_plus_b(controller_output, o2o_w, o2o_b)\n",
    "\n",
    "        state = {\n",
    "            'controller_state': controller_state,\n",
    "            'read_vector_list': read_vector_list,\n",
    "            'w_list': w_list,\n",
    "            'p_list': p_list,\n",
    "            'M': M\n",
    "        }\n",
    "\n",
    "        self.step += 1\n",
    "        \n",
    "        return NTM_output, state\n",
    "\n",
    "    \n",
    "    \n",
    "    def addressing(self, k, beta, g, s, gamma, prev_M, prev_w):\n",
    "        \"\"\"\n",
    "        We will implement two different addressing mechanisms,\n",
    "        \n",
    "        1. Content-based\n",
    "        2. Location-based\n",
    "            2.1. Interpolation\n",
    "            2.2. Convolutional Shift\n",
    "            2.3. Sharpening\n",
    "        \"\"\"\n",
    "       \n",
    "        \n",
    "        #1. Content based addressing\n",
    "        \n",
    "        #contoller returns a key vector K that is compared with the each row in the memory M using cosine similarity\n",
    "        k = tf.expand_dims(k, axis=2)\n",
    "        inner_product = tf.matmul(prev_M, k)\n",
    "        \n",
    "        k_norm = tf.sqrt(tf.reduce_sum(tf.square(k), axis=1, keepdims=True))\n",
    "        M_norm = tf.sqrt(tf.reduce_sum(tf.square(prev_M), axis=2, keepdims=True))\n",
    "        norm_product = M_norm * k_norm\n",
    "        \n",
    "        #Compute cosine similarity\n",
    "        K = tf.squeeze(inner_product / (norm_product + 1e-8))                     \n",
    "    \n",
    "        #now, we produce the normalized weight vector based on the similairty and the key strength (beta)\n",
    "        #beta is used for adjusting the precision on the head focus\n",
    "        \n",
    "        K_amplified = tf.exp(tf.expand_dims(beta, axis=1) * K)\n",
    "        w_c = K_amplified / tf.reduce_sum(K_amplified, axis=1, keepdims=True)       \n",
    "    \n",
    "\n",
    "        if self.addressing_mode == 'content':                                   \n",
    "            return w_c\n",
    "\n",
    "        \n",
    "        #2. Location based addressing\n",
    "        \n",
    "        #location based addressing involves three other steps,\n",
    "        #2.1 Interpolation\n",
    "        #2.2 Convolutional Shift\n",
    "        #2.3 Sharpening\n",
    "\n",
    "        \n",
    "        #2.1 Interpolation \n",
    "        \n",
    "        #The first step in the location based addressing is the interpolation.\n",
    "        #It is used to decide whether we should use the weights we obtained at the previous time step  or\n",
    "        #we should use the weights obtained through content based addressing. \n",
    "        #But how do we decide that? We use a new scalar parameter which is used for determining\n",
    "        #which weights we should use. \n",
    "\n",
    "        \n",
    "        g = tf.expand_dims(g, axis=1)\n",
    "        w_g = g * w_c + (1 - g) * prev_w                                       \n",
    "        \n",
    "        \n",
    "        #2.2 Convolutional Shift\n",
    "        #The next step is called convolution shit. It is used for moving the head position.\n",
    "        #i.e it is used for shifting the focus from one location to the another. \n",
    "        \n",
    "        s = tf.concat([s[:, :self.shift_range + 1],\n",
    "                       tf.zeros([s.get_shape()[0], self.memory_size - (self.shift_range * 2 + 1)]),\n",
    "                       s[:, -self.shift_range:]], axis=1)\n",
    "        \n",
    "        t = tf.concat([tf.reverse(s, axis=[1]), tf.reverse(s, axis=[1])], axis=1)\n",
    "        \n",
    "        s_matrix = tf.stack(\n",
    "            [t[:, self.memory_size - i - 1:self.memory_size * 2 - i - 1] for i in range(self.memory_size)],\n",
    "            axis=1\n",
    "        )\n",
    "        \n",
    "        w_ = tf.reduce_sum(tf.expand_dims(w_g, axis=1) * s_matrix, axis=2)      # eq (8)\n",
    "        \n",
    "        \n",
    "\n",
    "        \n",
    "        #2.3 Sharpening\n",
    "                \n",
    "        #The final step is sharpening. As a result of the convolutional shift, the weights will not be sharp.\n",
    "        #i.e because of the shift, weights focused at a single location will be dispersed into other locations.\n",
    "        #To mitigate this effect, we perform sharpening. \n",
    "        #We use another parameter called  gamma to perform sharpening and it can be expressed as,\n",
    "\n",
    "        \n",
    "        w_sharpen = tf.pow(w_, tf.expand_dims(gamma, axis=1))\n",
    "        w = w_sharpen / tf.reduce_sum(w_sharpen, axis=1, keepdims=True)\n",
    "\n",
    "        return w\n",
    "\n",
    "    \n",
    "    \n",
    "    #next we define the function called zero state for initializing all the states - \n",
    "    #controller state, read vector, weights and memory\n",
    "    \n",
    "    def zero_state(self, batch_size, dtype):\n",
    "        def expand(x, dim, N):\n",
    "            return tf.concat([tf.expand_dims(x, dim) for _ in range(N)], axis=dim)\n",
    "\n",
    "        with tf.variable_scope('init', reuse=self.reuse):\n",
    "            state = {\n",
    "                'controller_state': expand(tf.tanh(tf.get_variable('init_state', self.rnn_size,\n",
    "                                            initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))),\n",
    "                                  dim=0, N=batch_size),\n",
    "                \n",
    "                'read_vector_list': [expand(tf.nn.softmax(tf.get_variable('init_r_%d' % i, [self.memory_vector_dim],\n",
    "                                            initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))),\n",
    "                                  dim=0, N=batch_size)\n",
    "                           for i in range(self.read_head_num)],\n",
    "                \n",
    "                'w_list': [expand(tf.nn.softmax(tf.get_variable('init_w_%d' % i, [self.memory_size],\n",
    "                                            initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))),\n",
    "                                  dim=0, N=batch_size) if self.addressing_mode == 'content_and_loaction'\n",
    "                           else tf.zeros([batch_size, self.memory_size])\n",
    "                           for i in range(self.read_head_num + self.write_head_num)],\n",
    "                \n",
    "                'M': expand(tf.tanh(tf.get_variable('init_M', [self.memory_size, self.memory_vector_dim],\n",
    "                                            initializer=tf.random_normal_initializer(mean=0.0, stddev=0.5))),\n",
    "                                  dim=0, N=batch_size)\n",
    "            }\n",
    "            return state\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we define a function called generate random sequence which will generate random sequence of length seq_length and we will feed this seqence to the NTM input for the copy task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def generate_random_strings(batch_size, seq_length, vector_dim):\n",
    "    return np.random.randint(0, 2, size=[batch_size, seq_length, vector_dim]).astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create NTMCopyModel for performing the whole copy task. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NTMCopyModel():\n",
    "    \n",
    "    def __init__(self, args, seq_length, reuse=False):\n",
    "        \n",
    "        #input sequence\n",
    "        self.x = tf.placeholder(name='x', dtype=tf.float32, shape=[args.batch_size, seq_length, args.vector_dim])\n",
    "        \n",
    "        #output sequence\n",
    "        self.y = self.x\n",
    "        \n",
    "        #end of the sequence\n",
    "        eof = np.zeros([args.batch_size, args.vector_dim + 1])\n",
    "        eof[:, args.vector_dim] = np.ones([args.batch_size])\n",
    "        eof = tf.constant(eof, dtype=tf.float32)\n",
    "        zero = tf.constant(np.zeros([args.batch_size, args.vector_dim + 1]), dtype=tf.float32)\n",
    "        \n",
    "        if args.model == 'LSTM':\n",
    "            def rnn_cell(rnn_size):\n",
    "                return tf.nn.rnn_cell.BasicLSTMCell(rnn_size, reuse=reuse)\n",
    "            cell = tf.nn.rnn_cell.MultiRNNCell([rnn_cell(args.rnn_size) for _ in range(args.rnn_num_layers)])\n",
    "            \n",
    "            \n",
    "        elif args.model == 'NTM':\n",
    "            cell = NTMCell(args.rnn_size, args.memory_size, args.memory_vector_dim, 1, 1,\n",
    "                                    addressing_mode='content_and_location',\n",
    "                                    reuse=reuse,\n",
    "                                    output_dim=args.vector_dim)\n",
    "        \n",
    "        #initialize all the states\n",
    "        state = cell.zero_state(args.batch_size, tf.float32)\n",
    "        \n",
    "        self.state_list = [state]\n",
    "        \n",
    "        for t in range(seq_length):\n",
    "            output, state = cell(tf.concat([self.x[:, t, :], np.zeros([args.batch_size, 1])], axis=1), state)\n",
    "            self.state_list.append(state)\n",
    "        \n",
    "        #get the output and states\n",
    "        output, state = cell(eof, state)\n",
    "        self.state_list.append(state)\n",
    "\n",
    "        self.o = []\n",
    "        \n",
    "        for t in range(seq_length):\n",
    "            output, state = cell(zero, state)\n",
    "            self.o.append(output[:, 0:args.vector_dim])\n",
    "            self.state_list.append(state)\n",
    "        \n",
    "        self.o = tf.sigmoid(tf.transpose(self.o, perm=[1, 0, 2]))\n",
    "\n",
    "        eps = 1e-8\n",
    "        \n",
    "        #calculate loss as cross entropy loss\n",
    "        self.copy_loss = -tf.reduce_mean(self.y * tf.log(self.o + eps) + (1 - self.y) * tf.log(1 - self.o + eps))\n",
    "        \n",
    "        #optimize using RMS prop optimizer\n",
    "        with tf.variable_scope('optimizer', reuse=reuse):\n",
    "            self.optimizer = tf.train.RMSPropOptimizer(learning_rate=args.learning_rate, momentum=0.9, decay=0.95)\n",
    "            gvs = self.optimizer.compute_gradients(self.copy_loss)\n",
    "            capped_gvs = [(tf.clip_by_value(grad, -10., 10.), var) for grad, var in gvs]\n",
    "            self.train_op = self.optimizer.apply_gradients(capped_gvs)\n",
    "                        \n",
    "        self.copy_loss_summary = tf.summary.scalar('copy_loss_%d' % seq_length, self.copy_loss)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--mode', default=\"train\")\n",
    "parser.add_argument('--restore_training', default=False)\n",
    "parser.add_argument('--test_seq_length', type=int, default=5)\n",
    "parser.add_argument('--model', default=\"NTM\")\n",
    "parser.add_argument('--rnn_size', default=16)\n",
    "parser.add_argument('--rnn_num_layers', default=3)\n",
    "parser.add_argument('--max_seq_length', default=5)\n",
    "parser.add_argument('--memory_size', default=16)\n",
    "parser.add_argument('--memory_vector_dim', default=5)\n",
    "parser.add_argument('--batch_size', default=5)\n",
    "parser.add_argument('--vector_dim', default=8)\n",
    "parser.add_argument('--shift_range', default=1)\n",
    "parser.add_argument('--num_epoches', default=100)\n",
    "parser.add_argument('--learning_rate', default=1e-4)\n",
    "parser.add_argument('--save_dir', default= os.getcwd())\n",
    "parser.add_argument('--tensorboard_dir', default=os.getcwd())\n",
    "args = parser.parse_args(args = [])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's train the NTM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train(args):\n",
    "    model_list = [NTMCopyModel(args, 1)]\n",
    "    for seq_length in range(2, args.max_seq_length + 1):\n",
    "        model_list.append(NTMCopyModel(args, seq_length, reuse=True))\n",
    "\n",
    "        \n",
    "    with tf.Session() as sess:\n",
    "        if args.restore_training:\n",
    "            saver = tf.train.Saver()\n",
    "            ckpt = tf.train.get_checkpoint_state(args.save_dir + '/' + args.model)\n",
    "            saver.restore(sess, ckpt.model_checkpoint_path)\n",
    "        else:\n",
    "            saver = tf.train.Saver(tf.global_variables())\n",
    "            tf.global_variables_initializer().run()\n",
    "        \n",
    "        #initialize summary writer for visualizing in tensorboard\n",
    "        train_writer = tf.summary.FileWriter(args.tensorboard_dir, sess.graph)\n",
    "        plt.ion()\n",
    "        plt.show()\n",
    "    \n",
    "        for b in range(args.num_epoches):\n",
    "            \n",
    "            #initialize the sequence length\n",
    "            seq_length = np.random.randint(1, args.max_seq_length + 1)\n",
    "            model = model_list[seq_length - 1]\n",
    "            \n",
    "            #generate our random input sequence as an input\n",
    "            x = generate_random_strings(args.batch_size, seq_length, args.vector_dim)\n",
    "            \n",
    "            #feed our input to the model\n",
    "            feed_dict = {model.x: x}\n",
    "           \n",
    "            if b % 100 == 0:        \n",
    "                p = 0              \n",
    "                print(\"First training batch sample\",x[p, :, :])\n",
    "                \n",
    "                #compute model output\n",
    "                print(\"Model output\",sess.run(model.o, feed_dict=feed_dict)[p, :, :])\n",
    "                state_list = sess.run(model.state_list, feed_dict=feed_dict)\n",
    "                \n",
    "                \n",
    "                if args.model == 'NTM':\n",
    "                    w_plot = []\n",
    "                    M_plot = np.concatenate([state['M'][p, :, :] for state in state_list])\n",
    "                    for state in state_list:\n",
    "                        w_plot.append(np.concatenate([state['w_list'][0][p, :], state['w_list'][1][p, :]]))\n",
    "                    \n",
    "                    #plot the weight matrix to see the attention\n",
    "                    plt.imshow(w_plot, interpolation='nearest', cmap='gray')\n",
    "                    plt.draw()\n",
    "                    plt.pause(0.001)\n",
    "                \n",
    "                #compute loss\n",
    "                copy_loss = sess.run(model.copy_loss, feed_dict=feed_dict)\n",
    "                \n",
    "                #write to summary\n",
    "                merged_summary = sess.run(model.copy_loss_summary, feed_dict=feed_dict)\n",
    "                train_writer.add_summary(merged_summary, b)\n",
    "                \n",
    "                print('batches %d, loss %g' % (b, copy_loss))\n",
    "            else:                   \n",
    "                sess.run(model.train_op, feed_dict=feed_dict)\n",
    "                \n",
    "            #save the model\n",
    "            if b % 5000 == 0 and b > 0:\n",
    "                saver.save(sess, args.save_dir + '/' + args.model + '/model.tfmodel', global_step=b)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-2-c55831a50e07>:104: calling expand_dims (from tensorflow.python.ops.array_ops) with dim is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use the `axis` argument instead\n",
      "('First training batch sample', array([[1., 0., 1., 1., 0., 1., 0., 1.],\n",
      "       [0., 0., 0., 1., 1., 0., 0., 1.],\n",
      "       [0., 1., 1., 1., 1., 1., 1., 1.],\n",
      "       [1., 0., 1., 0., 0., 0., 1., 0.]], dtype=float32))\n",
      "('Model output', array([[0.41622064, 0.40011194, 0.8426822 , 0.18761651, 0.52241313,\n",
      "        0.4941342 , 0.58150166, 0.27654198],\n",
      "       [0.39070413, 0.50569034, 0.6326411 , 0.23488827, 0.6771403 ,\n",
      "        0.53548634, 0.5629458 , 0.23213845],\n",
      "       [0.52530485, 0.5213146 , 0.76477057, 0.24487178, 0.6035095 ,\n",
      "        0.38190985, 0.49055338, 0.36256135],\n",
      "       [0.4900252 , 0.6046341 , 0.6669247 , 0.27418157, 0.54876125,\n",
      "        0.38766316, 0.5006444 , 0.33060235]], dtype=float32))\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW4AAACLCAYAAABBXHMKAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAADYVJREFUeJzt3VuMXfV1x/HfzzO+jgffGQVsahcMfbAqE0FE2mBFCUZpWoVWiiKQXMV9cR8aStpKrdsXUktIVZVG6RMVTahSFTKNwGmtIppYKqjkxbXxpb5hY1MTbBxfGGw84MsYrz6cjTSx53jWhrPn+D/6fiTLM2eW1/7/539m+cw++7+XI0IAgHJM6fYAAAD1ULgBoDAUbgAoDIUbAApD4QaAwlC4AaAwFG4AKAyFGwAKQ+EGgML0NpHUNtsxAaCmiHAmjlfcAFCYVOG2/SXbB2wfsr2+6UEBANrzeDeZst0j6aCk1ZKOStoq6ZGI2Hedf8OpEgCoqZOnSj4j6VBEvBERlyQNSnrokwwOAPDxZQr3rZLeGvX50eqxX2J7ne1ttrd1anAAgGt17KqSiHhK0lMSp0oAoEmZV9zHJC0Z9fni6jEAQBdkCvdWScttL7M9TdLDkjY1OywAQDvjniqJiMu2vyHpJ5J6JD0dEXsbHxkAYEzjXg74sZLaYY9/Vctdd92Vznno0KFU3LRp09I5Z82alYpbsGBBOufhw4dTcZnvz0emT5+eihsZGUnn3LBhQyru5ZdfTufcunVrKu7ChQvpnNnYnp6eVNzFixfTx25Cdt1nz56dznnu3LlU3NKlS9M5jxw5kopbs2ZNOuezzz6bihscHEznXLt2bSpuxowZ6ZzZ59ydd96Zzrlz585UHDsnAWCSonADQGEo3ABQGAo3ABSGwg0AhaFwA0BhKNwAUBgKNwAUhsINAIWhcANAYbq65b3Otu+BgYFUXF9fXzrnu+++m4qrs036gw8+SMdmTZ06NRW3atWqdM7bb789FbdpU/5+YnPnzk3F9fbm7ya8b1/bRku/5PLly6m4Os+5rOz6SFL2563OlvdLly6l4vr7+9M5s7dZqHP7gtWrV6fiXnzxxY7nfOGFF9I5s8/js2fPpnNmttwPDQ1pZGSELe8AMBmNW7htL7H9ku19tvfafmwiBgYAGFvmd9bLkv4sIrbb7pf0qu3N12sWDABozrivuCPieERsrz4+J2m/xug5CQCYGLV6TtpeKuluSVvG+No6Ses6MioAQFvpwm17tqTnJX0zIt67+us0CwaAiZG6qsT2VLWK9jMRsbHZIQEAridzVYklfV/S/oj4TvNDAgBcT+YV929K+n1JX7C9s/rz5YbHBQBoI9Pl/WeSam83y+wQq7NrM7tD6/3330/nvHLlSipueHg4nbMJ2Z152UbFUr556ZkzZ9I5H3300VTc22+/nc6Z3bV64MCBVNyUKfk9Z9nYOk2aZ86cmYqrs1s3+zyus2v09OnTqbgHHnggnXPjxtxZ1ieeeCKdc/369am4e++9N51zy5Zrrr0YU7bZuJT7fmbXUWLnJAAUh8INAIWhcANAYSjcAFAYCjcAFIbCDQCFoXADQGEo3ABQGAo3ABSGwg0AhWmsWXCncz744IOpuM2bN6dzNjH3JmS3vNfZMvvhhx+m4lauXJnOuWPHjlRcthmrVK8ha0adRsVz5sxJxdVpEJ1pGitJixYtSuccGhpKxd10003pnNkGxNkmzZK0YsWKVNzu3bvTOZcsWZKKy94SQcr/vL333jV3t25r4cKF48acPn2aZsEAMFmlC7ftHts7bP9HkwMCAFxfnVfcj6nVbxIA0EXZDjiLJf22pO81OxwAwHiyr7i/K+nPJbV998v2OtvbbG/ryMgAAGPKtC77HUknI+LV68VFxFMRcU9E3NOx0QEArpFtXfYV20ckDarVwuxfGh0VAKCtcQt3RPxlRCyOiKWSHpb0XxGxpvGRAQDGxHXcAFCYru6crNNsc8GCBam4OrvD9u7dm47Nyu7Mq/N9v/nmm1Nx2QavktTT05OKGxgYSOfcsGFDKm5wcDCdM7sTNruDr85zbunSpam4U6dOpXPOmzcvFVdnN2b2uVSnUXK26fby5cvTOV9//fVUXPb7LuV3RPb396dzvvPOO6m4bONnKdfsfGRkRFeuXGHnJABMRhRuACgMhRsACkPhBoDCULgBoDAUbgAoDIUbAApD4QaAwlC4AaAwFG4AKEy+c2oDzp8/3/HYO+64I51z37596disbMPeOk1rs01W62zhz26TrtMM9pVXXknF1dkinm1qnDV9+vR0bHbrc53t1NkGs9lGxZJ05syZVFxfX186Zzb2xIkT6ZzZ21G8+eabHc85MjKSzpm9LUH2tgBS7nlX52eNV9wAUJhs67K5tp+z/Zrt/bY/2/TAAABjy/6+/veS/jMivmp7mqT8LdYAAB01buG2PUfSKklrJSkiLkm61OywAADtZE6VLJN0StI/2d5h+3u2r3nngmbBADAxMoW7V9KnJT0ZEXdLel/S+quDaBYMABMjU7iPSjoaEVuqz59Tq5ADALog0yz4F5Lesn1X9dAXJXX+AmgAQEr2qpJHJT1TXVHyhqQ/aG5IAIDrSRXuiNgpqavnroeHh1Nxx48fT+ecOnVqKs5O9e+UlN/9VCfna6+9lopbtmxZOufhw4dTcXV2hx07diwVd99996VzZpvBnjt3LhVXZz7ZNcruhpTyTa+zz00pv1u3TiPto0ePpuKy85HyuywXLlyYznny5MlU3IwZM9I5L13KXTSX3WEpSUNDQ+nYDHZOAkBhKNwAUBgKNwAUhsINAIWhcANAYSjcAFAYCjcAFIbCDQCFoXADQGEo3ABQmK42C842rJXy21AHBgbSOQ8ePJiKmzlzZjpntrltdj6StGjRolTcbbfdls6ZHef+/fvTOffs2ZOOzaqzRT2jznMuu6X54sWL6Zy33HJLKq7OrRtmzco1pKrTpDn7nK+TM7vtPNukWcrfGuDChQvpnFOm5F7P1rnVQZ1bXGTwihsACpNtFvwntvfa3mP7h7bzd2wBAHTUuIXb9q2S/ljSPRGxQlKPpIebHhgAYGzZUyW9kmba7lWrw/vbzQ0JAHA9mQ44xyR9W9LPJR2XdDYifnp1HM2CAWBiZE6VzJP0kFrd3m+R1Gd7zdVxNAsGgImROVXygKT/i4hTETEiaaOk32h2WACAdjKF++eS7rM9y62LEb8oKX9xLwCgozLnuLdIek7Sdkm7q3/zVMPjAgC0kW0W/Likxxsey3Vldx7V2RmX3SFVZ9dT9vjZY0v5hqj3339/OueuXbtScXV2eGYbJc+fPz+ds6enJxWXbZibzSdJZ8+eTcXNmTMnnTO707BOY99ujrPOruLz58+n4vr7+9M5s02ip0+fns6Z3Qnb25vfeJ7JWat2pSMBADcECjcAFIbCDQCFoXADQGEo3ABQGAo3ABSGwg0AhaFwA0BhKNwAUBgKNwAUxnW2WaaT2qckvXnVwwslne74wbqH+dz4JtucmM+N75PM6VciItUZvJHCPeaB7G2T6V7dzOfGN9nmxHxufBM1J06VAEBhKNwAUJiJLNyT7R7ezOfGN9nmxHxufBMypwk7xw0A6AxOlQBAYRov3La/ZPuA7UO21zd9vIlg+4jt3bZ32t7W7fHUZftp2ydt7xn12Hzbm22/Xv09r5tjrKPNfL5l+1i1Rjttf7mbY6zD9hLbL9neZ3uv7ceqx0teo3ZzKnKdbM+w/T+2d1Xz+evq8WW2t1T17l9tT2vk+E2eKrHdI+mgpNWSjkraKumRiNjX2EEngO0jku6JiCKvQbW9StKwpH+OiBXVY38raSgi/qb6D3ZeRPxFN8eZ1WY+35I0HBHf7ubYPg7bn5L0qYjYbrtf0quSflfSWpW7Ru3m9DUVuE5V4/S+iBi2PVXSzyQ9JulPJW2MiEHb/yBpV0Q82enjN/2K+zOSDkXEGxFxSdKgpIcaPibGERH/LWnoqocfkvSD6uMfqPVDVYQ28ylWRByPiO3Vx+ck7Zd0q8peo3ZzKlK0DFefTq3+hKQvqNVcXWpwjZou3LdKemvU50dV8GKNEpJ+avtV2+u6PZgOGYiI49XHv5A00M3BdMg3bP9vdSqlmNMKo9leKuluSVs0SdboqjlJha6T7R7bOyWdlLRZ0mFJZyLio67ZjdU73pz8eD4XEZ+W9FuS/qj6VX3SiNb5s9IvN3pS0u2SVko6Lunvujuc+mzPlvS8pG9GxHujv1bqGo0xp2LXKSI+jIiVkhardXbh1ybq2E0X7mOSloz6fHH1WNEi4lj190lJP1Zr0Up3ojoP+dH5yJNdHs8nEhEnqh+sK5L+UYWtUXXe9HlJz0TExurhotdorDmVvk6SFBFnJL0k6bOS5trurb7UWL1runBvlbS8eqd1mqSHJW1q+JiNst1Xvbki232SHpS05/r/qgibJH29+vjrkv69i2P5xD4qcJXfU0FrVL3x9X1J+yPiO6O+VOwatZtTqetke5HtudXHM9W6AGO/WgX8q1VYY2vU+Aac6vKe70rqkfR0RDzR6AEbZvtX1XqVLUm9kp4tbU62fyjp82rdyeyEpMcl/ZukH0m6Ta07O34tIop4w6/NfD6v1q/fIemIpD8cdX74hmb7c5JekbRb0pXq4b9S65xwqWvUbk6PqMB1sv3rar352KPWC+AfRcSGqj4MSpovaYekNRFxsePHZ+ckAJSFNycBoDAUbgAoDIUbAApD4QaAwlC4AaAwFG4AKAyFGwAKQ+EGgML8P6GUgNBd5CeAAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batches 0, loss 0.786682\n"
     ]
    }
   ],
   "source": [
    "train(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Credits to some parts of the code used in this section goes to [this]( https://github.com/MarkPKCollier/NeuralTuringMachine) github repo. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
